from learning_to_adapt.spaces.base import Space
import numpy as np


class Product(Space):
    """
    表示一个由多个空间组成的乘积空间。
    """

    def __init__(self, *components):
        if isinstance(components[0], (list, tuple)):
            assert len(components) == 1
            components = components[0]
        self._components = tuple(components)  # 由多个空间组成的元组
        dtypes = [c.new_tensor_variable("tmp", extra_dims=0).dtype for c in components]
        if len(dtypes) > 0 and hasattr(dtypes[0], "as_numpy_dtype"):
            dtypes = [d.as_numpy_dtype for d in dtypes]
        self._common_dtype = np.core.numerictypes.find_common_type([], dtypes)  # 找到所有空间的最小公倍数

    def sample(self):
        return tuple(x.sample() for x in self._components)  # 对各个空间进行采样并返回

    @property
    def components(self):
        return self._components  # 返回组成此乘积空间的各个空间

    def contains(self, x):
        return isinstance(x, tuple) and all(c.contains(xi) for c, xi in zip(self._components, x))  # 判断给定元素是否在此乘积空间中

    @property
    def flat_dim(self):
        return int(np.sum([c.flat_dim for c in self._components]))  # 返回此乘积空间扁平化后的向量维度

    def flatten(self, x):
        return np.concatenate([c.flatten(xi) for c, xi in zip(self._components, x)])  # 对各个空间的元素进行扁平化操作并返回

    def flatten_n(self, xs):
        xs_regrouped = [[x[i] for x in xs] for i in range(len(xs[0]))]  # 对给定批量元素进行重新分组
        flat_regrouped = [c.flatten_n(xi) for c, xi in zip(self.components, xs_regrouped)]  # 对重新分组后的元素进行扁平化操作
        return np.concatenate(flat_regrouped, axis=-1)  # 返回扁平化后的批量元素

    def unflatten(self, x):
        dims = [c.flat_dim for c in self._components]  # 获取各个空间的扁平化向量维度
        flat_xs = np.split(x, np.cumsum(dims)[:-1])  # 将扁平化向量按照维度切割
        return tuple(c.unflatten(xi) for c, xi in zip(self._components, flat_xs))  # 对各个切割后的向量进行还原操作并返回

    def unflatten_n(self, xs):
        dims = [c.flat_dim for c in self._components]  # 获取各个空间的扁平化向量维度
        flat_xs = np.split(xs, np.cumsum(dims)[:-1], axis=-1)  # 将扁平化向量按照维度切割
        unflat_xs = [c.unflatten_n(xi) for c, xi in zip(self.components, flat_xs)]  # 对各个切割后的向量进行还原操作
        unflat_xs_grouped = list(zip(*unflat_xs))  # 将还原后的向量重新分组
        return unflat_xs_grouped  # 返回重新分组后的向量

    @property
    def default_value(self):
        return tuple([x.default_value for x in self.components])  # 返回各个空间的默认值构成的元组

    def __eq__(self, other):
        if not isinstance(other, Product):
            return False
        return tuple(self.components) == tuple(other.components)  # 判断两个乘积空间是否相等

    def __hash__(self):
        return hash(tuple(self.components))  # 返回乘积空间对象的哈希值
